//
// Created by jg on 2021/1/7.
//


#include "vis_rviz/vis_rviz.h"

void VIS_RVIZ::init(ros::NodeHandle &nh) {
    nh.param("vis/vis_traj_width", vis_traj_width_, 0.15);
    nh.param("vis/is_proj_cube",   is_proj_cube_,   true);

    vis_map_pub_      = nh.advertise<visualization_msgs::Marker>("/vis_map",         1000);
    vis_traj_pub_     = nh.advertise<visualization_msgs::Marker>("/trajectory_vis",     1);
    vis_rob_pub_      = nh.advertise<visualization_msgs::Marker>("/robot_pos",         10);
    vis_target_pub_   = nh.advertise<visualization_msgs::Marker>("/target_pos",        20);
    vis_pr_pub_       = nh.advertise<visualization_msgs::Marker>("/reference_traj",     1);
    vis_path_pub_     = nh.advertise<visualization_msgs::MarkerArray>("/vis_path",     10);
    vis_corridor_pub_ = nh.advertise<visualization_msgs::MarkerArray>("/vis_corridor",  1);

    bspline_pub_  = nh.advertise<visualization_msgs::Marker>("/planning_vis/bspline_traj", 20);
    traj_pub_     = nh.advertise<visualization_msgs::Marker>("/planning_vis/bspline_traj", 20);
    pubs_.push_back(bspline_pub_);
    pubs_.push_back(traj_pub_);
}

void VIS_RVIZ::setEnv(MAP_EDT::Ptr &map_edt, Bernstein::Ptr& bezier) {
    map_edt_ = map_edt;
    bezier_  = bezier;
}

void VIS_RVIZ::visMap(VIS_MAP_TYPE map_type) {
    setMapMarker(map_type);

    map_marker_.header.frame_id = "odom";
    map_marker_.header.stamp = ros::Time::now();
    map_marker_.id = 0;

    if(map_type == GLOBALMAP)
        map_edt_->setGlobalMapMarker(map_marker_);
    else if(map_type == LOCALMAP)
        map_edt_->setLocalMapMarker(map_marker_);

    vis_map_pub_.publish(map_marker_);
}

void VIS_RVIZ::setMapMarker(VIS_MAP_TYPE map_type) {
    if(map_type == GLOBALMAP){
        /*decide from which view we can see the marker*/
        map_marker_.ns = std::string("world");
        /*set marker type*/
        map_marker_.type = visualization_msgs::Marker::SPHERE_LIST;
        /*decide if the marker will be enlarge*/
        map_marker_.scale.x = 0.1;
        map_marker_.scale.y = 0.1;
        map_marker_.scale.z = 0.1;
        map_marker_.pose.orientation.w = 1;
        map_marker_.pose.orientation.x = 0;
        map_marker_.pose.orientation.y = 0;
        map_marker_.pose.orientation.z = 0;
        /*decide the color of the marker*/
        map_marker_.color.a = 1.0; // Don't forget to set the alpha!
        map_marker_.color.r = 1.0;
        map_marker_.color.g = 1.0;
        map_marker_.color.b = 1.0;
        /*set marker action*/
        map_marker_.action = visualization_msgs::Marker::ADD;
    } else if(map_type == LOCALMAP){
        /*decide from which view we can see the marker*/
        map_marker_.ns = "local";
        /*set marker type*/
        map_marker_.type = visualization_msgs::Marker::SPHERE_LIST;
        /*decide if the marker will be enlarge*/
        map_marker_.scale.x = 0.15;
        map_marker_.scale.y = 0.15;
        map_marker_.scale.z = 0.15;
        map_marker_.pose.orientation.w = 1;
        map_marker_.pose.orientation.x = 0;
        map_marker_.pose.orientation.y = 0;
        map_marker_.pose.orientation.z = 0;
        /*decide the color of the marker*/
        map_marker_.color.a = 1.0; // Don't forget to set the alpha!
        map_marker_.color.r = 1.0;
        map_marker_.color.g = 0.0;
        map_marker_.color.b = 0.0;
        /*set marker action*/
        map_marker_.action = visualization_msgs::Marker::ADD;
    }
}

void VIS_RVIZ::visRobotPos(Eigen::Vector3d &pt, double yaw) {
    visualization_msgs::Marker pt_vis;

    pt_vis.header.stamp       = ros::Time::now();
    pt_vis.header.frame_id    = "odom";

    pt_vis.ns = "robot/position";
    pt_vis.id = 0;
    pt_vis.type = visualization_msgs::Marker::ARROW;

    pt_vis.action = visualization_msgs::Marker::DELETE;
    vis_rob_pub_.publish(pt_vis);
//    _checkTraj_vis_pub.publish(traj_vis);
//    _stopTraj_vis_pub.publish(traj_vis);

    pt_vis.action = visualization_msgs::Marker::ADD;
    const float _lenght = 0.8;
    pt_vis.scale.x = 0.3;
    pt_vis.scale.y = 0.4;
    pt_vis.scale.z = _lenght;
//    pt_vis.pose.orientation.x = 0.0;
//    pt_vis.pose.orientation.y = 0.0;
//    pt_vis.pose.orientation.z = 0.0;
//    pt_vis.pose.orientation.w = 1.0;
    pt_vis.color.r = 1.0;
    pt_vis.color.g = 1.0;
    pt_vis.color.b = 0.0;
    pt_vis.color.a = 1.0;

    geometry_msgs::Point p1, p2;
    p1.x = pt(0) - 0.5*_lenght * cos(yaw);
    p1.y = pt(1) - 0.5*_lenght * sin(yaw);
    p1.z = pt(2);
    p2.x = pt(0) + 0.5*_lenght * cos(yaw);
    p2.y = pt(1) + 0.5*_lenght * sin(yaw);
    p2.z = pt(2);
    pt_vis.points.push_back(p1);
    pt_vis.points.push_back(p2);

    vis_rob_pub_.publish(pt_vis);

    robot_traj_.push_back(pt);
    displaySphereList(robot_traj_, 0.08, Eigen::Vector4d(0.0, 0.7, 1.0, 1), 1);
}

void VIS_RVIZ::resetVisRviz() {
    target_marker_.points.clear();
    vis_target_pub_.publish(target_marker_);

    pr_marker_.points.clear();
    vis_pr_pub_.publish(pr_marker_);

    robot_traj_.clear();
}

void VIS_RVIZ::visTargetPt(const Eigen::Vector3d &pt) {

    target_marker_.header.stamp       = ros::Time::now();
    target_marker_.header.frame_id    = "odom";

    target_marker_.ns = "target/position";
    target_marker_.id = 0;
    target_marker_.type = visualization_msgs::Marker::SPHERE_LIST;

//    pt_vis.action = visualization_msgs::Marker::DELETE;
//    vis_target_pub_.publish(pt_vis);

    target_marker_.action = visualization_msgs::Marker::ADD;
    target_marker_.scale.x = 0.3;
    target_marker_.scale.y = 0.3;
    target_marker_.scale.z = 0.3;
    target_marker_.pose.orientation.x = 0.0;
    target_marker_.pose.orientation.y = 0.0;
    target_marker_.pose.orientation.z = 0.0;
    target_marker_.pose.orientation.w = 1.0;
    target_marker_.color.r = 1.0;
    target_marker_.color.g = 0.0;
    target_marker_.color.b = 0.0;
    target_marker_.color.a = 1.0;

    geometry_msgs::Point point;
    point.x = pt[0];
    point.y = pt[1];
    point.z = pt[2];
    target_marker_.points.push_back(point);

    vis_target_pub_.publish(target_marker_);
}

void VIS_RVIZ::visPath(const vector<Eigen::Vector3d>& path) {

    for(auto& mk: path_vis_.markers)
        mk.action = visualization_msgs::Marker::DELETE;

    vis_path_pub_.publish(path_vis_);
    path_vis_.markers.clear();
    if(path.empty()) return;

    visualization_msgs::Marker mk;
    mk.header.frame_id = "odom";
    mk.header.stamp = ros::Time::now();
    mk.ns = "b_traj/fast_marching_path";
    mk.type = visualization_msgs::Marker::CUBE;
    mk.action = visualization_msgs::Marker::ADD;

    mk.pose.orientation.x = 0.0;
    mk.pose.orientation.y = 0.0;
    mk.pose.orientation.z = 0.0;
    mk.pose.orientation.w = 1.0;
    mk.color.a = 0.6;
    mk.color.r = 0.8;
    mk.color.g = 0.8;
    mk.color.b = 0.0;

    int idx = 0;
    double resolution = map_edt_->getResolution();
    for(int i = 0; i < int(path.size()); i++){
        mk.id = idx;

        mk.pose.position.x = path[i](0);
        mk.pose.position.y = path[i](1);
        mk.pose.position.z = path[i](2);

        mk.scale.x = resolution;
        mk.scale.y = resolution;
        mk.scale.z = resolution;

        idx ++;
        path_vis_.markers.push_back(mk);
    }

    vis_path_pub_.publish(path_vis_);
}

void VIS_RVIZ::visCorridor(const vector<Cube> &corridor) {
    for(auto & mk: cube_vis_.markers)
        mk.action = visualization_msgs::Marker::DELETE;

    vis_corridor_pub_.publish(cube_vis_);
    cube_vis_.markers.clear();
    if(corridor.empty()) return;

    visualization_msgs::Marker mk;
    mk.header.frame_id = "odom";
    mk.header.stamp = ros::Time::now();
    mk.ns = "corridor";
    mk.type = visualization_msgs::Marker::CUBE;
    mk.action = visualization_msgs::Marker::ADD;

    mk.pose.orientation.x = 0.0;
    mk.pose.orientation.y = 0.0;
    mk.pose.orientation.z = 0.0;
    mk.pose.orientation.w = 1.0;

    mk.color.a = 0.4;
    mk.color.r = 1.0;
    mk.color.g = 1.0;
    mk.color.b = 1.0;

    int idx = 0;
    for(int i = 0; i < int(corridor.size()); i++)
    {
        mk.id = idx;

        mk.pose.position.x = (corridor[i].vertex(0, 0) + corridor[i].vertex(3, 0) ) / 2.0;
        mk.pose.position.y = (corridor[i].vertex(0, 1) + corridor[i].vertex(1, 1) ) / 2.0;

        if(is_proj_cube_)
            mk.pose.position.z = 0.0;
        else
            mk.pose.position.z = (corridor[i].vertex(0, 2) + corridor[i].vertex(4, 2) ) / 2.0;

        mk.scale.x = (corridor[i].vertex(0, 0) - corridor[i].vertex(3, 0) );
        mk.scale.y = (corridor[i].vertex(1, 1) - corridor[i].vertex(0, 1) );

        if(is_proj_cube_)
            mk.scale.z = 0.05;
        else
            mk.scale.z = (corridor[i].vertex(0, 2) - corridor[i].vertex(4, 2) );

        idx ++;
        cube_vis_.markers.push_back(mk);
    }

    vis_corridor_pub_.publish(cube_vis_);
}

void VIS_RVIZ::visBezierTraj(const Eigen::MatrixXd &poly_coeff, const Eigen::VectorXd &time) {
    traj_vis_.header.stamp       = ros::Time::now();
    traj_vis_.header.frame_id    = "odom";

    traj_vis_.ns = "trajectory/bezier";
    traj_vis_.id = 0;
    traj_vis_.type = visualization_msgs::Marker::SPHERE_LIST;

    traj_vis_.action = visualization_msgs::Marker::DELETE;
    vis_traj_pub_.publish(traj_vis_);
    traj_vis_.points.clear();
//    checkTraj_vis_pub_.publish(traj_vis);
//    stopTraj_vis_pub_.publish(traj_vis);

    traj_vis_.action = visualization_msgs::Marker::ADD;
    traj_vis_.scale.x = vis_traj_width_;
    traj_vis_.scale.y = vis_traj_width_;
    traj_vis_.scale.z = vis_traj_width_;
    traj_vis_.pose.orientation.x = 0.0;
    traj_vis_.pose.orientation.y = 0.0;
    traj_vis_.pose.orientation.z = 0.0;
    traj_vis_.pose.orientation.w = 1.0;
    traj_vis_.color.r = 0.0;
    traj_vis_.color.g = 1.0;
    traj_vis_.color.b = 0.7;
    traj_vis_.color.a = 1.0;

    double traj_len = 0.0;
    int count = 0;
    Eigen::Vector3d cur, pre;
    cur.setZero();
    pre.setZero();

    Eigen::Vector3d rob_pos;
    geometry_msgs::Point pt;

    int segment_num  = poly_coeff.rows();
    for(int i = 0; i < segment_num; i++ ){
        for (double t = 0.0; t < 1.0; t += 0.05 / time(i), count += 1){
            rob_pos = bezier_->getPosFromBezier( poly_coeff, t, i );
            cur(0) = pt.x = time(i) * rob_pos(0);
            cur(1) = pt.y = time(i) * rob_pos(1);
            cur(2) = pt.z = time(i) * rob_pos(2);
            traj_vis_.points.push_back(pt);

            if (count) traj_len += (pre - cur).norm();
            pre = cur;
        }
    }

    ROS_INFO("[GENERATOR] The length of the trajectory; %.3lfm.", traj_len);
    vis_traj_pub_.publish(traj_vis_);
}

void VIS_RVIZ::visReferenTraj(const Eigen::Vector3d &pr) {

    pr_marker_.header.stamp       = ros::Time::now();
    pr_marker_.header.frame_id    = "odom";

    pr_marker_.ns = "reference/traj";
    pr_marker_.id = 0;
    pr_marker_.type = visualization_msgs::Marker::SPHERE_LIST;

    pr_marker_.action = visualization_msgs::Marker::ADD;
    pr_marker_.scale.x = vis_traj_width_;
    pr_marker_.scale.y = vis_traj_width_;
    pr_marker_.scale.z = vis_traj_width_;
    pr_marker_.pose.orientation.x = 0.0;
    pr_marker_.pose.orientation.y = 0.0;
    pr_marker_.pose.orientation.z = 0.0;
    pr_marker_.pose.orientation.w = 1.0;
    pr_marker_.color.r = 0.0;
    pr_marker_.color.g = 1.0;
    pr_marker_.color.b = 0.7;
    pr_marker_.color.a = 1.0;

    geometry_msgs::Point point;
    point.x = pr[0];
    point.y = pr[1];
    point.z = pr[2];
    pr_marker_.points.push_back(point);

    vis_pr_pub_.publish(pr_marker_);
}

void VIS_RVIZ::visBsplineTraj(NonUniformBspline &bspline, double size, const Eigen::Vector4d &color, bool show_ctrl_pts,
                              double size2, const Eigen::Vector4d &color2, int id1, int id2) {
    if (bspline.getControlPoint().size() == 0) return;

    vector<Eigen::Vector3d> traj_pts;
    double tm, tmp;
    bspline.getTimeSpan(tm, tmp);

    for(double t = tm; t <= tmp; t += 0.01) {
        Eigen::Vector3d pt = bspline.evaluateDeBoor(t);
        traj_pts.push_back(pt);
    }
    displaySphereList(traj_pts, size, color, BSPLINE + id1 % 100);

    if(!show_ctrl_pts) return;
    Eigen::MatrixXd         ctrl_pts = bspline.getControlPoint();
    vector<Eigen::Vector3d> ctp;

    for (int i = 0; i < int(ctrl_pts.rows()); ++i) {
        Eigen::Vector3d pt = ctrl_pts.row(i).transpose();
        ctp.push_back(pt);
    }
    displaySphereList(ctp, size2, color2, BSPLINE_CTRL_PT + id2 % 100);

}

void VIS_RVIZ::displaySphereList(const vector<Eigen::Vector3d>& list, double resolution,
                                              const Eigen::Vector4d& color, int id, int pub_id) {
    visualization_msgs::Marker mk;
    mk.header.frame_id = "odom";
    mk.header.stamp    = ros::Time::now();
    mk.type            = visualization_msgs::Marker::SPHERE_LIST;
    mk.action          = visualization_msgs::Marker::DELETE;
    mk.id              = id;
    pubs_[pub_id].publish(mk);

    mk.action             = visualization_msgs::Marker::ADD;
    mk.pose.orientation.x = 0.0;
    mk.pose.orientation.y = 0.0;
    mk.pose.orientation.z = 0.0;
    mk.pose.orientation.w = 1.0;

    mk.color.r = color(0);
    mk.color.g = color(1);
    mk.color.b = color(2);
    mk.color.a = color(3);

    mk.scale.x = resolution;
    mk.scale.y = resolution;
    mk.scale.z = resolution;

    geometry_msgs::Point pt;
    for (int i = 0; i < int(list.size()); i++) {
        pt.x = list[i](0);
        pt.y = list[i](1);
        pt.z = list[i](2);
        mk.points.push_back(pt);
    }
    pubs_[pub_id].publish(mk);
    ros::Duration(0.001).sleep();
}
